from Crypto.Cipher import AES
from Crypto.Util import Counter
import dbEngine

class CipherEngine:
    def __init__(self,key,db):
        self.key = key
        self.db = db
    def __initCounter(self,filename=None, IV=None):
        if not IV:
            IV = self.db.getIV(filename)
            if not IV:
                import math, random
                rand = random.Random()
                IV = rand.randint(1,int(rand.random() * math.pow(10,7)))
        return IV, Counter.new(64, "\xDE\xAD\xBE\xEF", "\xAA\xAA\xAA\xAA", initial_value=IV, little_endian=False)
    
    def encrypt(self,data,filename=None,IV=None):
        IV, counter = self.__initCounter(filename,IV)
        aes = AES.new(self.key,AES.MODE_CTR,counter=counter)
        return IV, aes.encrypt(data)
    
    def decrypt(self,data,filename=None,IV=None):
        IV, counter = self.__initCounter(filename,IV)
        aes = AES.new(self.key,AES.MODE_CTR,counter=counter)
        return aes.decrypt(data)